import click
from learn_flask.extensions import db
from learn_flask.models import User, Movie


def register_commands(app):
    """
    注册自定义Flask命令
    
    Args:
        app (Flask): Flask应用实例
    """
    
    @app.cli.command('init-db')
    @click.option('--drop', is_flag=True, help='Create after drop.')
    def init_database(drop):
        """
        初始化数据库命令
        
        Args:
            drop (bool): 是否先删除现有数据库
            
        通过命令行执行：flask init-db 或 flask init-db --drop
        """
        if drop:
            # 如果指定了--drop选项，则删除所有表
            db.drop_all()
            click.echo('Dropped all tables.')
        # 创建所有表
        db.create_all()
        click.echo('Initialized database.')

    @app.cli.command()
    def forge():
        """
        生成测试数据命令
        
        通过命令行执行：flask forge
        """
        # 先清空数据库并重新创建表
        db.drop_all()
        db.create_all()

        # 预设用户信息
        name = 'Grey Li'
        # 预设电影数据
        movies = [
            {'title': 'My Neighbor Totoro', 'year': '1988'},
            {'title': 'Dead Poets Society', 'year': '1989'},
            {'title': 'A Perfect World', 'year': '1993'},
            {'title': 'Leon', 'year': '1994'},
            {'title': 'Mahjong', 'year': '1996'},
            {'title': 'Swallowtail Butterfly', 'year': '1996'},
            {'title': 'King of Comedy', 'year': '1999'},
            {'title': 'Devils on the Doorstep', 'year': '1999'},
            {'title': 'WALL-E', 'year': '2008'},
            {'title': 'The Pork of Music', 'year': '2012'},
        ]

        # 创建用户并设置密码
        user = User(name=name, username='grey')
        user.set_password('123')
        db.session.add(user)
        
        # 添加所有预设电影
        for m in movies:
            movie = Movie(title=m['title'], year=m['year'])
            db.session.add(movie)

        # 提交事务
        db.session.commit()
        click.echo('Done.')